/*
 * $AIST_Release: 0.9.0 $
 * Copyright 2011 Information Technology Research Institute, National
 * Institute of Advanced Industrial Science and Technology
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *    http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package parser.analyzer

import scala.collection.immutable.{Set => MSet}
import scala.collection.mutable.{Map => MMap}
import scala.collection.mutable.ListBuffer
import scala.util.parsing.input.Position
import scala.util.parsing.input.NoPosition
import scala.collection.JavaConversions._
import com.google.protobuf.Descriptors._
import com.google.protobuf.DescriptorProtos._
import aggr.Aggregators
import aggr.Aggregators.aggrMap
import aggr.Aggregators.AggregatorFactory
import aggr.{SerializedTypes => ST}
import util.PBUtils
import parser.Nodes._
import parser.DataTypes._
import parser.TreeWalker
import parser.Messages
import parser.Results._
import util.Tools
import util.FileUtils
import TableSerializationInfoMaker.mkKeyTupleType
import TableSerializationInfoMaker.mkRealValueTypeFromTuple
import TableSerializationInfoMaker.mkArrayTuple

/**
 * Module for semantic analysis(Main stage).
 */
object MainAnalyzer {

  private def registerConverter(vt: VariableTable, id: Int, dataType: MessageType): Result[Unit] = {
    val t2b = FunctionType(BytesType, List(TypeReference(id)))
    val b2t = FunctionType(TypeReference(id), List(BytesType))

    (fromOption(dataType.pos, vt.registerFunction(BytesType.name, t2b)) ~ 
     fromOption(dataType.pos, vt.registerFunction(dataType.desc.getName, b2t))).mapConst(Unit)
  }

  /**
   * Result of analyze().
   */
  case class State(ast: List[Stmt])

  /**
   * Analyzes AST.
   */
  def analyze(
    ast: List[Stmt], 
    typeMap: Map[Int, DataType], 
    globals: Iterable[(String, DataType)]): Result[State] = 
  {
    val variableTable = {
      val v = new VariableTable(typeMap, globals)
      val rs = typeMap.toList.flatMap { case (id, dataType) =>
        dataType match {
          case m: MessageType => List(registerConverter(v, id, m))
          case _              => Nil
        }
      }
      mergeList(rs).map(_ => v)
    }

    variableTable.flatMap { vt =>
      val ma = new MainAnalyzer(typeMap, vt)

      ma.analyze(ast).map { ast =>
        State(ast)
      }
    }
  }
}

private class MainAnalyzer(val typeMap: Map[Int, DataType], variableTable: VariableTable) extends ExprAnalyzer {

  override def findConverter(from: DataType, to: DataType): Option[Expr] = variableTable.findConverter(from, to)
  override def resolveVariable(v: Variable): Result[Expr] = {
    (variableTable.findVariable(v.name) match {
      case Some(VariableInfo(dataType, pos, VariableTypeGlobal)) => 
        Success(GlobalVariable(v.name, dataType, pos))
      case Some(VariableInfo(dataType, pos, VariableTypeStatic)) => 
        Success(StaticVariable(v.name, dataType, pos))
      case Some(VariableInfo(dataType, pos, VariableTypeLocal)) => 
        if (isStatic) failure(v.pos, "\"%s\" is local variable. Local variables cannot be referred from initializer of static variable." , v.name)
        else          Success(LocalVariable(v.name, RefType(dataType), pos))
      case None => 
        failure(v.pos, "\"%s\" is undeclared" , v.name)
    }).map { x => x.setPos(v.pos) }
  }

  private val retTypeStack = new ListBuffer[DataType]

  var isStatic = false
  private def staticExpr[T](f: => T): T = {
    isStatic = true
    try     f
    finally isStatic = false;
  }

  override def handleFuncLit(funcLit: FuncLit): Result[(List[Arg], List[Stmt])] = {
    val FuncLit(retType, args, stmts) = funcLit
    inScope {
      retTypeStack.append(retType)
      try {
        val r: List[Result[Arg]] = args.map { label =>
          label match {
            case Label("", dataType) =>  Success(UnusedArg(dataType))
            case Label(name, dataType) => 
              variableTable.registerLocalVar(name, dataType, false) match {
                case Left(pos)  => Success(UsedArg(name, pos, dataType))
                case Right(msg) => failure(label.pos, msg)
              }
          }
        }
        mergeList(r).flatMap { args => 
          def checkReturnStmt(stmts: List[Stmt]): Result[Unit] = {
            stmts.lastOption match {
              case Some(s: ReturnStmt)           => Success(Unit)
              case Some(Scope(s))                => checkReturnStmt(List(s))
              case Some(StmtList(ss))            => checkReturnStmt(ss)
              case Some(IfStmt(_, ts, Some(es))) => checkReturnStmt(List(ts)).flatMap { _ => checkReturnStmt(List(es)) }
              case x => failure(x.map(_.pos).getOrElse(funcLit.pos), "The tail of function is not return statement.")
            }
          }

          analyze(stmts).flatMap { stmts => 
            checkReturnStmt(stmts).mapConst((args, stmts))
          }
        }
      } finally {
        retTypeStack.trimEnd(1)
      }
    }
  }

  private val enumMap: Map[String, EnumType] = typeMap.flatMap { case (_, e) =>
    e match {
      case e: EnumType => Map(e.desc.getName -> e)
      case _ => Nil
    }
  }

  override def getEnumValue(e: String, elem: String): Result[Expr] = {
    enumMap.get(e) match {
      case Some(dataType) => 
        val value = dataType.desc.findValueByName(elem)
        if (value == null) {
          Failure(Nil) // No message.
        } else {
          Success(EnumExpr(value.getNumber, dataType))
        }
      case None => Failure(Nil) // No message.
    }
  }

  private def inScope[T](f: => T): T = {
    variableTable.pushBlock()
    try f
    finally variableTable.popBlock()
  }

  private var nLoop = 0
  private def inLoop[T](f: => T): T = {
    nLoop += 1
    try f
    finally nLoop -=1
  }

  def analyze(ast: List[Stmt]) = {
    TreeWalker.Mapper(classOf[Stmt]).setFollower {
      f => {
        case s: Scope =>
          inScope(f(s.stmt)).map { case s: Stmt => Scope(s) }
        case s@LoopStmt(cond, stmt, update) => inLoop {
          val condR = analyzeCond(cond) 
          val updateR = processOpt(update)(analyzeExpr)

          (condR ~ f(stmt) ~ updateR).map { case cond ~ (stmt: Stmt) ~ update =>
            LoopStmt(cond, stmt, update).setPos(s.pos)
          }
        }
        case s:WhenStmt =>
          analyzeWhenStmt(s, f)

        case s: AnalyzedWhenStmt => error("Unexpect AnalyzedWhenStmt.")
      }
    }.setPostprocess {
      case s: ExprStmt => analyzeExpr(s.expr).map { ExprStmt(_).setPos(s.pos)}
      case s@DefVarStmt(mod, names, dataType, initalValue) =>
        mod match {
          case Some("static") => analyzeStaticDefVar(s.pos, names, dataType, initalValue)
          case _              => 
            if (names.length > 1 && initalValue.isDefined) {
              failure(s.pos, "When two variables or more are declared at the same time, it is not possible to initialize them.")
            } else {
              val initR = initalValue match {
                case Some(x) => analyzeInitilizer(x, dataType) 
                case None    => Success(UndefLit(dataType))
              }
              val posR = mergeList(names.map { n =>
                variableTable.registerLocalVar(n.name, dataType, false) match {
                  case Right(msg) => failure(n.pos, msg)
                  case Left(pos)  => Success((n.name, pos))
                }
              })

              (initR ~ posR).map { case iv ~ poss =>
                StmtList(poss.toList.map { case (name, pos) => VarInit(name, pos, iv) }).setPos(s.pos)
              }
            }
        }
      case _: DefTableStmt => error("Invalid Statement")
      case s:DefTableStmtWithSerializationInfo => 
        analyzeDefTableStmtWithSerializationInfo(s)
      case s: EmitStmt => analyzeEmitStmt(s)
      case s: IfStmt =>
        analyzeCond(s.cond).map { cond =>
          s.copy(cond = cond).setPos(s.pos)
        }
      case s: DefStaticVarStmt => error("Does not exist")
      case s: BreakStmt =>
        if (nLoop == 0) failure(s.pos, "break statement is not within loop.")
        else            Success(s)
      case s: ContinueStmt =>
        if (nLoop == 0) failure(s.pos, "continue statement is not within loop.")
        else            Success(s)
      case s: ReturnStmt => 
        analyzeExpr(s.value).flatMap { expr =>
          retTypeStack.lastOption match {
            case Some(t) => checkType(expr, t).mapConst(ReturnStmt(expr))
            case None    => failure(s.pos, "Unexpect \"return\", here is not in function.")
          }
        }
    }.map(ast).map{_.asInstanceOf[List[Stmt]]}
  }

  private def analyzeWhenStmt(s: WhenStmt, f: Any => Result[Any]): Result[Any] = {
    val WhenStmt(defvars, cond, stmt) = s

    def registerWhenVars(vars: List[(Identifier, WhenMod, DataType)]): Result[List[(String, Int)]] = {
      val poss = vars.map { case (name, _, dataType) =>
        variableTable.registerLocalVar(name.name, dataType, false) match {
          case Left(pos)  => Left( (name.name, pos) )
          case Right(msg) => Right( Messages.error(name.pos, msg) )
        }
      } 
      val (posList, errors) = poss.partition{_.isLeft}

      errors match {
        case Nil => Success(posList.flatMap{_.left.toOption})
        case _   => Failure(errors.flatMap{_.right.toOption})
      }
    }

    def bindWhenVar(varSet: Set[Int], expr: Expr): Map[Int, Expr] = {
      val map = MMap[Int, Expr]()
      TreeWalker.walkTree(classOf[Expr], expr) {
        case ContainerAccessOp(container, ValueExpr(LocalVariable(_, _, vpos)), _) => 
          if (varSet.contains(vpos) && !map.contains(vpos)) {
            map += (vpos -> container)
          }
          Success(Unit)
      }
      return Map() ++ map
    }

    val vars = defvars.flatMap { case WhenDefVar(names, mod, dataType) => names.map { (_, mod, dataType) } }

    inScope {
      registerWhenVars(vars.filter{_._2 != AllMod}).flatMap { posList1 =>
        inScope (
          // Can access variables with "all" modifier in only condition expression.
          registerWhenVars(vars.filter{_._2 == AllMod}).flatMap { posList => 
            analyzeCond(cond).map { cond => (cond, posList) }
          }
        ).flatMap { case (cond, posList2) =>
          val varSet = vars.map { _._1 }

          val posMap = Map() ++ posList1 ++ posList2
          val conMap = bindWhenVar(Set() ++ posMap.values, cond)

          varSet.filter(x => !conMap.contains(posMap(x.name))).toList match {
            case Nil => f(stmt).map { case stmt: Stmt => 
              AnalyzedWhenStmt(vars.map { x => 
                  val p = posMap(x._1.name)
                  WhenVar(x._1.name, p, conMap(p), x._2) 
                }, cond, stmt)
            }
            case unusedVars => 
              Failure(unusedVars.map(x => Messages.error(x.pos, "Does not bind to container: %s", x.name) ))
          }
        }
      }
    }
  }

  private def analyzeDefTableStmtWithSerializationInfo(s: DefTableStmtWithSerializationInfo): Result[Stmt] = {
    val DefTableStmtWithSerializationInfo(names, tInfo, sInfo) = s 
    val TableInfo(aggr, args, idxs, conType, opts) = tInfo
    val TableSerializationInfo(sKey, sArgs, sRealVal) = sInfo

    mergeList(names.map { case vn =>
      val (sVal, sOpts) = splitValOpts(opts.nonEmpty, sRealVal)
      val checkAggrR = aggrMap.get(aggr.name) match {
        case Some(af) => 
          af.check(sKey, sVal, sArgs, sOpts) match {
            case None      => Success(Unit)
            case Some(Nil) => failure(s.pos, "Invalid table format \"%s\".".format(aggr.name))
            case Some(x)   => Failure(x.map{ Messages.error(s.pos, "%s", _) } )
          }
        case None => failure(aggr.pos, "Unkown aggregator \"%s\".".format(aggr.name))
      }

      val optsMap = opts.map { case TableOption(id, label) => (id, label.dataType) }
      for (_     <- checkAggrR;
           pargs <- staticExpr(processOpt(args)(analyzeExpr));
           _     <- processOpt(pargs) { checkType(_, IntType) };
           rslt  <- {
            val dataType = TableType(idxs.map{_.dataType}, conType.dataType, optsMap)
            variableTable.registerLocalVar(vn.name, dataType, true) match {
              case Left(id) => 
                Success(TableInit(id, vn.name, aggr.name, pargs.toList, dataType, sInfo))
              case Right(msg) => failure(vn.pos, msg)
            }
           }
      ) yield rslt
    }).map{StmtList(_)}
  }

  private def splitValOpts(hasOpt: Boolean, real: ST.SerializedType): (ST.SerializedType, Map[String, ST.SerializedType]) = {
    if (hasOpt) {
      assert(real.isInstanceOf[ST.MessageType])
      val msg = real.asInstanceOf[ST.MessageType]
      val flds = msg.descriptor.getFields()

      val (v, o) = flds.partition(_.getName == "data__")

      assert(v.length == 1)
      assert(o.length > 0)

      (ST.getType(v(0)), Map() ++ o.map{x => (x.getName, ST.getType(x))})
    } else {
      (real, Map())
    }
  }

  private def analyzeEmitStmt(s: EmitStmt): Result[Stmt] = {
    val EmitStmt(table, keys, expr, opts) = s
        
    def checkKeyTypes(keys: List[Expr], expectKeys: List[DataType]): Result[Unit] = {
      val keyTypes = keys.map{_.dataType}
      if (keyTypes.corresponds(expectKeys){matchType(typeMap, _, _)}) {
        Success(Unit)
      } else {
        failure(table.pos, "Keys type mismatch:\n\tfound: %s\n\trequired: ", 
          "[" + keyTypes.mkString(", ") + "]",
          "[" + expectKeys.mkString(", ") + "]")
      }
    }

    def checkEmitOptions(opts: List[EmitOption], expect: List[(String, DataType)]): Result[Unit] = {
      val expectMap = Map() ++ expect
      val optsMap0 = Map[String, List[EmitOption]]().withDefaultValue(List())
      val optsMap = opts.foldLeft(optsMap0) { case (optsMapN, opt) =>
        optsMapN.updated(opt.id, optsMapN(opt.id) :+ opt)
      }
      val keys = (expectMap.keySet ++ optsMap.keySet).toList

      val results = keys.map{ key => 
        (expectMap.get(key), optsMap(key)) match {
          case (Some(expectType), x::xs) => 
            checkType(x.value, expectType) ~ mergeList(xs.map(x => failure(x.pos, "%s options has already used.", x.id)))
          case (Some(_), Nil) => failure(s.pos, "\"%s\" is required, but not exist", key)
          case (None, xs)     => mergeList(xs.map(x => failure(x.pos, "Unexpect option \"%s\"", x.id) ))
        }
      }
      return mergeList(results).map(_ => Unit)
    }

    def byteFuncCall(src: Expr): Expr = {
      val s = stripRef(src.dataType) match {
        case x: ArrayType => TupleExpr(List(stripRefExpr(src)), mkArrayTuple(x))
        case _ => src
      }
      FuncCall(BuildInFunction(BytesType.name, FunctionType(BytesType, List(stripRef(s.dataType)))), List(s))
    }
    
    val tableState  = analyzeExpr(table)
    val keysState   = mergeList(keys.map(analyzeExpr))
    val targetState = analyzeExpr(expr)
    val optsState   = opts.map { x => analyzeExpr(x.value).map { i => EmitOption(x.id, i).setPos(x.pos)} }

    (tableState ~ keysState ~ targetState ~ mergeList(optsState)).flatMap { case table ~ keys ~ target ~ opts =>
      stripRef(table.dataType) match {
        case TableType(expectKeys, dataType, expectOpts) =>
          for (_ <- checkType(target, dataType);
               _ <- checkKeyTypes(keys, expectKeys);
               _ <- checkEmitOptions(opts, expectOpts)) yield {
            val k = keys match {
              case Nil => UndefLit(EmptyType)
              case List(x) => byteFuncCall(x)
              case x => 
                byteFuncCall(TupleExpr(x, mkKeyTupleType(expectKeys).asInstanceOf[TupleType]))
            }
            val v = expectOpts match {
              case Nil => byteFuncCall(target)
              case x => 
                byteFuncCall(TupleExpr(target::(opts.map{_.value}), 
                mkRealValueTypeFromTuple(stripRef(target.dataType), expectOpts).asInstanceOf[TupleType]))
            }
            AnalyzedEmitStmt(table, k, v)
          }
        case x => failure(table.pos, "Emit target is not table.(%s)", x)
      }
    }
  }

  private def analyzeCond(expr: Expr): Result[Expr] = {
    // Condition type is int.
    for (cond <- analyzeExpr(expr).map(stripRefExpr);
         _    <- checkType(cond, IntType)) yield cond
  }

  private def analyzeStaticDefVar(pos: Position, names: List[Identifier], dataType: DataType, initalValue: Option[Initializer]): Result[Stmt] = {

    def analyzeSomeInitializer(d: DataType): Result[Expr] = initalValue match {
      case None    => failure(pos, "Static variable must be initialized,")
      case Some(i) => analyzeInitilizer(i, d)
    }

    def getVarName(): Result[Identifier] = names match {
      case Nil        => error("Empty variable names")
      case List(name) => Success(name)
      case _          => failure(pos, "When two variables or more are declared at the same time, it is not possible to initialize them.")
    }

    staticExpr {
      for {
        n <- getVarName()
        i <- analyzeSomeInitializer(dataType)
        s <- {
          variableTable.registerLocalVar(n.name, dataType, true) match {
            case Right(msg) => failure(n.pos, msg)
            case Left(id) => Success(DefStaticVarStmt(id, n.name, i))
          }
        }
      } yield s
    }
  }
}
